{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SUknhuHKyc-E"
   },
   "source": [
    "# <center>OpenAI agent pattern: routing</center>\n",
    "\n",
    "A starter guide for building an agent loop using the `openai-agents` library.\n",
    "\n",
    "This pattern uses routing to choose which specialized agent to use for a specific sub-task. The specialized agent attempts to complete the sub-task and return a result.\n",
    "\n",
    "In the following example, we'll build an agent which creates a portfolio of stocks and ETFs based on a user's investment strategy.\n",
    "1.  **Router Agent:** Chooses which worker to use based on the user's investment strategy.\n",
    "2.  **Research Agent:** Searches the web for information about stocks and ETFs that could support the user's investment strategy.\n",
    "3.  **Question Answering Agent:** Answers questions about investing like Warren Buffett."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n69HR7eJswNt"
   },
   "source": [
    "### Install Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install base libraries for OpenAI\n",
    "!pip install -q openai openai-agents\n",
    "\n",
    "# Install optional libraries for OpenInference/OpenTelemetry tracing\n",
    "!pip install -q arize-phoenix-otel openinference-instrumentation-openai-agents openinference-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jQnyEnJisyn3"
   },
   "source": [
    "### Setup Keys\n",
    "\n",
    "Add your OpenAI API key to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Copy your Phoenix `API_KEY` from your settings page at [app.phoenix.arize.com](https://app.phoenix.arize.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kfid5cE99yN5"
   },
   "source": [
    "### Setup Tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(\n",
    "    project_name=\"openai-agents\",\n",
    "    auto_instrument=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from textwrap import dedent\n",
    "\n",
    "from agents import Agent, Runner, TResponseInputItem, WebSearchTool\n",
    "from agents.model_settings import ModelSettings\n",
    "\n",
    "qa_agent = Agent(\n",
    "    name=\"Investing Q&A Agent\",\n",
    "    instructions=dedent(\"\"\"You are Warren Buffett. You are answering questions about investing.\"\"\"),\n",
    "    model=\"gpt-4.1\",\n",
    ")\n",
    "\n",
    "research_agent = Agent(\n",
    "    name=\"Financial Search Agent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a research assistant specializing in financial topics. Given a stock ticker, use web search to retrieve up‑to‑date context and produce a short summary of at most 50 words. Focus on key numbers, events, or quotes that will be useful to a financial analyst.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    "    tools=[WebSearchTool()],\n",
    "    model_settings=ModelSettings(tool_choice=\"required\", parallel_tool_calls=True),\n",
    ")\n",
    "\n",
    "orchestrator_agent = Agent(\n",
    "    name=\"Routing Agent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a senior financial analyst. Your task is to handoff to the appropriate agent or tool.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    "    handoffs=[\n",
    "        research_agent,\n",
    "        qa_agent,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_items: list[TResponseInputItem] = []\n",
    "\n",
    "while True:\n",
    "    user_input = input(\"Enter your question: \")\n",
    "    if user_input == \"exit\":\n",
    "        break\n",
    "    input_item = {\"content\": user_input, \"role\": \"user\"}\n",
    "    input_items.append(input_item)\n",
    "    orchestrator = await Runner.run(orchestrator_agent, input_items)\n",
    "    orchestrator_output = orchestrator.final_output\n",
    "    pprint(orchestrator.last_agent)\n",
    "    pprint(orchestrator_output)\n",
    "    input_items = orchestrator.to_input_list()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
